from vanna.ZhipuAI.ZhipuAI_Chat import ZhipuAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore


class ZhipuAIVanna(ChromaDB_VectorStore, ZhipuAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        ZhipuAI_Chat.__init__(self, config=config)


class MyVanna:
    def __init__(self, api_key: str):
        vn = ZhipuAIVanna(

            config={"model": "glm-4", "api_key": api_key})  # chatglm2-6b

        vn.connect_to_mysql(host='10.240.53.162',
                            dbname='test',
                            user='rdsroot',
                            password='Geely@db20211206',
                            port=3306)
        vn.train(ddl="""
        CREATE TABLE IF NOT EXISTS users (
            id INT PRIMARY KEY COMMENT '用户ID'， 
            username VARCHAR(50) COMMENT '用户名'， 
            email VARCHAR(100) COMMENT '电子邮件'， 
            age INT COMMENT '年龄'， 
            gender VARCHAR(10) COMMENT '性别（男/女）'， 
            city VARCHAR(50) COMMENT '城市'
        ) COMMENT='用户信息表' CHARACTER SET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
        """)

        self.vanna = vn
